#include <cstdio>
#include <cuda.h>
#include <cuda_runtime.h>

// v0: naive版本
// latency: 3.835ms
// blockSize作为模板参数的效果主要用于静态shared memory的申请需要传入编译期常量指定大小（L10)
template <int blockSize>
__global__ void reduce_v0(float *d_in, float *d_out) {
    __shared__ float smem[blockSize];
    // 当前block中的thread id
    int tid = threadIdx.x;
    // 全局的thread id
    int gtid = blockIdx.x * blockDim.x + threadIdx.x;
    // load: 每个线程加载一个元素到shared mem对应位置
    smem[tid] = d_in[gtid];
    // 涉及到对shared memory的读写最好都加上__syncthreads
    __syncthreads();
    /*                                             工作tid
    s:1 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | 0   2   4   6   8   10    12    14    16
    s:2 0   2   4   6   8   10    12    14    16 | 0       4       8         12          16
    s:4 0       4       8         12          16 | 0               8                     16
    s:8 0               8                     16 | 0
        0
    */
    for (int s = 1; s < blockDim.x; s *= 2) {
        if (tid % (2 * s) == 0) {
            smem[tid] = smem[tid] + smem[tid + s];
        }
        __syncthreads();
    }

    // store: 哪里来回哪里去，把reduce结果写回显存
    if (tid == 0) {
        d_out[blockIdx.x] = smem[0];
    }
}

bool check_result(float *out, float groudtruth, int n) {
    float res = 0;
    for (int i = 0; i < n; i++) {
        res += out[i];
    }
    if (res != groudtruth) {
        return false;
    }
    return true;
}

int main() {
    float milliseconds = 0;
    const int N = 25600000;
    cudaSetDevice(0);
    cudaDeviceProp deviceProp;
    cudaGetDeviceProperties(&deviceProp, 0);
    const int blockSize = 256;
    int gridSize = std::min((N + 256 - 1) / 256, deviceProp.maxGridSize[0]);
    // int GridSize = 100000;
    float *a = (float *)malloc(N * sizeof(float));
    float *d_a;
    cudaMalloc((void **)&d_a, N * sizeof(float));

    float *out = (float *)malloc((gridSize) * sizeof(float));
    float *d_out;
    cudaMalloc((void **)&d_out, (gridSize) * sizeof(float));

    for (int i = 0; i < N; i++) {
        a[i] = 1.0f;
    }

    float groudtruth = N * 1.0f;

    cudaMemcpy(d_a, a, N * sizeof(float), cudaMemcpyHostToDevice);

    dim3 gridDim(gridSize);
    dim3 blockDim(blockSize);

    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    cudaEventRecord(start);
    reduce_v0<blockSize><<<gridDim, blockDim>>>(d_a, d_out);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&milliseconds, start, stop);

    cudaMemcpy(out, d_out, gridSize * sizeof(float), cudaMemcpyDeviceToHost);
    printf("allcated %d blocks, data counts are %d", gridSize, N);
    bool is_right = check_result(out, groudtruth, gridSize);
    if (is_right) {
        printf("the ans is right\n");
    } else {
        printf("the ans is wrong\n");
        // for(int i = 0; i < GridSize;i++){
        // printf("res per block : %lf ",out[i]);
        // }
        // printf("\n");
        printf("groudtruth is: %f \n", groudtruth);
    }
    printf("reduce_v0 latency = %f ms\n", milliseconds);

    cudaFree(d_a);
    cudaFree(d_out);
    free(a);
    free(out);
}
